# -*- coding: utf-8 -*-
import json

from conf.my_log import MyLog
from conf.neo4j_cnx import NeoCnxPool
from medical_kg.entity_type import entity_view
from medical_kg.prop_set_up import prop_view

logger = MyLog().get_logger()

neo_cnx = NeoCnxPool()


class MedKg:

    def __init__(self):
        self.rel_list = []
        self.node_list = []
        self.uid_list = []

    @staticmethod
    def query_rel_by_uid(uid):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (n:ICD11 {uid:$uid}) '
                'MATCH (n)-[r]->() '
                'RETURN DISTINCT type(r) AS type_r'
            )
            results = session.run(query, uid=uid)
            results = [record['type_r'] for record in results.data()]
            return results

    @staticmethod
    def query_next_nodes(uid, rel):
        with neo_cnx.get_session() as session:
            query = (
                        'MATCH (n1:ICD11 {uid:$uid}) '
                        'MATCH (n1)-[:%s]->(n2) '
                        'RETURN n2'
                    ) % rel
            results = session.run(query, uid=uid)
            results = [record['n2'] for record in results.data()]
            return results

    @staticmethod
    def query_property(uid, type):
        with neo_cnx.get_session() as session:
            if type == 'NODE':
                query = (
                    'MATCH (n:ICD11 {uid:$uid}) '
                    'RETURN '
                    'n.properties as properties'
                )
                results = session.run(query, uid=uid)
                for record in results:
                    return prop_view(json.loads(record['properties']))

    def query_oneself(self, uid):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (n:ICD11 {uid:$uid}) '
                'RETURN '
                'n'
            )
            results = session.run(query, uid=uid)

            for record in results:
                for i in ['n']:
                    labels = list(record[i].labels)
                    labels.remove('ICD11')
                    global_properties = dict(record[i].items())
                    node_uid = global_properties['uid']
                    properties = json.loads(global_properties['properties'])

                    if node_uid not in self.uid_list:
                        self.uid_list.append(node_uid)
                        self.node_list.append({'id': str(node_uid),
                                               'labels': entity_view(labels),
                                               'properties': prop_view(properties)})

    def query_disease_kg(self, uid, rel_type_list):
        self.query_oneself(uid)

        for rel_type in rel_type_list:
            # TODO: 回路查找图谱
            self.query_disease_kg_by_designated_rel(uid, rel_type)

        data = {'graph': {'nodes': self.node_list, 'relationships': self.rel_list}}

        return data

    def query_disease_kg_by_designated_rel(self, uid, rel_type):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (n1:ICD11 {uid:$uid}) '
                'MATCH (n1)-[r]->(n2) '
                'WHERE type(r) = $rel_type '
                'RETURN n1, n2, r'
            )
            results = session.run(query, uid=uid, rel_type=rel_type)

            for record in results:
                for i in ['n1', 'n2']:
                    labels = list(record[i].labels)
                    labels.remove('ICD11')

                    global_properties = dict(record[i].items())
                    node_uid = global_properties['uid']
                    properties = json.loads(global_properties['properties'])

                    if node_uid not in self.uid_list:
                        self.uid_list.append(node_uid)
                        self.node_list.append({'id': str(node_uid),
                                               'labels': entity_view(labels),
                                               'properties': prop_view(properties)})
                for i in ['r']:
                    global_properties = dict(record[i].items())

                    properties = {}
                    if 'properties' in global_properties:
                        properties = json.loads(global_properties['properties'])
                    rel_uid = global_properties['uid']
                    rel_type = record[i].type
                    start_uid = record[i].nodes[0].get('uid')
                    end_uid = record[i].nodes[1].get('uid')

                    if rel_uid not in self.uid_list:
                        self.uid_list.append(id)
                        self.rel_list.append({
                            'id': str(rel_uid),
                            'type': rel_type,
                            'startNode': start_uid,
                            'endNode': end_uid,
                            'properties': prop_view(properties, type='REL')})


if __name__ == '__main__':
    mkg_app = MedKg()
    data = mkg_app.query_disease_kg(uid='257068234', rel_type_list=['Parent'])
    print(json.dumps(data, ensure_ascii=False))
    mkg_app.query_next_nodes(uid='1533297008', rel='体征')
    # mkg_app.del_rel(uid1='1533297008', uid2='1323682030')
    neo_cnx.close()
